#include "stmt.h"
#include <assert.h>
#include <ctype.h>
#include <memory.h>
#include <string.h>
#include <stdlib.h>
#include "common.h"
#include <map>
#include <string>

using namespace stmt;

#ifndef WIN32
#define _strdup strdup
#endif

// #define PRINT_PROCESS

#if defined PRINT_PROCESS && defined DEBUG
#define dbg(x) x
#else
#define dbg(...)
#endif

#define EXPECT_TOKEN(stmt, token) \
do {\
    if (tokc.kind != token) {\
        mark_line();\
        expect_token(token);\
        delete stmt;\
        stmt = nullptr;\
        return nullptr;\
    }\
} while (0)

#define NEXT_ON_KIND(stmt, token) \
do {\
    if (tokc.kind != token) {\
        mark_line();\
        expect_token(token);\
        delete stmt;\
        stmt = nullptr;\
        return nullptr;\
    } else {\
        token_next();\
    }\
} while (0)

void do_print_statement(ast_node_t* node, int tab, bool s);

namespace stmt {
decl_t* parse_decl();
}

void block_stmt_t::push(ast_node_t* stmt)
{
    svs_push_ptr(&statements, stmt);
}

DEFINE_SV(stmt_ptr_t)

if_stmt_t* parse_if_stmt()
{
    auto* stmt = new if_stmt_t;
    NEXT_ON_KIND(stmt, TK_IF);
    NEXT_ON_KIND(stmt, TK_LPAREN);
    stmt->cond = (expr_t*)parse_expression();
    if (stmt->cond == nullptr) {
        delete stmt;
        return nullptr;
    }
    NEXT_ON_KIND(stmt, TK_RPAREN);
    stmt->then_stmt = (stmt_t*)parse_statement();
    if (stmt->then_stmt == nullptr) {
        delete stmt;
        return nullptr;
    }
    if (tokc.kind == TK_ELSE) {
        token_next();
        stmt->else_stmt = (stmt_t*)parse_statement();
        if (stmt->else_stmt == nullptr) {
            delete stmt;
            return nullptr;
        }
    }
    return stmt;
}

while_stmt_t* parse_while_stmt()
{
    auto* stmt = new while_stmt_t;
    NEXT_ON_KIND(stmt, TK_WHILE);
    NEXT_ON_KIND(stmt, TK_LPAREN);
    stmt->cond = (expr_t*)parse_expression();
    if (stmt->cond == nullptr) {
        delete stmt;
        return nullptr;
    }
    NEXT_ON_KIND(stmt, TK_RPAREN);
    stmt->stmt = (stmt_t*)parse_statement();
    if (stmt->stmt == nullptr) {
        delete stmt;
        return nullptr;
    }
    return stmt;
}

for_stmt_t* parse_for_stmt()
{
    auto* stmt = new for_stmt_t;
    NEXT_ON_KIND(stmt, TK_FOR);
    NEXT_ON_KIND(stmt, TK_LPAREN);
    stmt->init = (stmt_t*)parse_statement();
    stmt->cond = (stmt_t*)parse_statement();
    stmt->post = (stmt_t*)parse_statement();
    NEXT_ON_KIND(stmt, TK_RPAREN);
    stmt->stmt = (stmt_t*)parse_statement();
    if (stmt->stmt == nullptr) {
        delete stmt;
        return nullptr;
    }
    return stmt;
}

block_stmt_t* parse_block_stmt()
{
    auto* stmt = new block_stmt_t;
    NEXT_ON_KIND(stmt, TK_LBRACE);
    while (tokc.kind != TK_RBRACE) {
        auto sub = parse_node();
        if (sub == nullptr) {
            delete stmt;
            return nullptr;
        }
        stmt->push(sub);
    }
    NEXT_ON_KIND(stmt, TK_RBRACE);
    return stmt;
}

expr_stmt_t* parse_expr_stmt()
{
    auto* stmt = new expr_stmt_t;
    stmt->expr = (expr_t*)parse_expression();
    if (stmt->expr == nullptr) {
        delete stmt;
        return nullptr;
    }
    if (tokc.kind == TK_SEMICOLON)
        token_next();
    return stmt;
}

import_stmt_t* parse_import_stmt()
{
    if (tokc.kind == TK_IMPORT) {
        auto* stmt = new import_stmt_t;
        stmt->token = tokc;
        NEXT_ON_KIND(stmt, TK_IMPORT);
        if (tokc.kind == TK_SEMICOLON)
            token_next();
        return stmt;
    }
    return nullptr;
}

method_t* parse_method()
{
    if (tokc.kind != TK_DEF) {
        return nullptr;
    }
    auto* stmt = new method_t;
    int line = tokc.loc.line;

    // def
    NEXT_ON_KIND(stmt, TK_DEF);
    // function name
    EXPECT_TOKEN(stmt, TK_ID);
    stmt->name = _strdup(ref_to_string(&tokc.ss));
    token_next();
    // (
    EXPECT_TOKEN(stmt, TK_LPAREN);
    token_next();
    // parameters
    if (tokc.kind == TK_ID) {
        do {
            if (tokc.kind == TK_COMMA) {
                token_next();
            }
            param_t* param = (param_t*)parse_decl();
            if (param == nullptr) {
                mark_line();
                do_error(&tokc.loc, "illegal param");
                delete stmt;
                return nullptr;
            }
            stmt->push(param);
        } while (tokc.kind == TK_COMMA);
    }
    // )
    EXPECT_TOKEN(stmt, TK_RPAREN);
    token_next();
    // {
    EXPECT_TOKEN(stmt, TK_LBRACE);
    token_next();
    // statements
    while (tokc.kind != TK_RBRACE && tokc.kind != TK_NONE) {
        auto sub = parse_statement();
        if (sub == nullptr) {
            delete stmt;
            return nullptr;
        }
        stmt->push((stmt_t*)sub);
    }
    // }
    EXPECT_TOKEN(stmt, TK_RBRACE);
    token_next();
    if (tokc.kind == TK_SEMICOLON)
        token_next();

    return stmt;
}

return_stmt_t* parse_return_stmt()
{
    auto* stmt = new return_stmt_t;
    EXPECT_TOKEN(stmt, TK_RETURN);
    token_next();
    stmt->expr = (expr_t*)parse_expression();
    if (stmt->expr == nullptr) {
        delete stmt;
        return nullptr;
    }
    if (tokc.kind == TK_SEMICOLON)
        token_next();
    return stmt;
}

global_stmt_t* parse_global_stmt()
{
    auto* stmt = new global_stmt_t;
    NEXT_ON_KIND(stmt, TK_GLOBAL);
    EXPECT_TOKEN(stmt, TK_ID);
    stmt->expr = (expr_t*)parse_expression();
    if (stmt->expr == nullptr) {
        delete stmt;
        return nullptr;
    }
    if (tokc.kind == TK_SEMICOLON)
        token_next();
    return stmt;
}

stmt_t* parse_break_stmt()
{
    auto* stmt = new stmt_t;
    NEXT_ON_KIND(stmt, TK_BREAK);
    stmt->kind = NK_BreakStatement;
    stmt->subkind = STATEMENT_BREAK;
    token_next();
    if (tokc.kind == TK_SEMICOLON)
        token_next();
    return stmt;
}

namespace stmt {

bool do_parse_decl(decl_t* decl)
{
	assert(tokc.kind == TK_ID);
    decl->token = tokc;
    decl->loc = tokc.loc;
    decl->decl = parse_primary_expression();
    assert(decl->decl);
    decl->name = _strdup(ref_to_string(&decl->decl->token.ss));
    if (tokc.kind == TK_ASSIGN) {
        token_next();
        auto cur = tokc;
        decl->init = parse_cond_expression();
        if (decl->init == nullptr) {
            mark_line();
            do_error(&cur.loc, "expect expression here");
            return false;
        }
    }
    if (tokc.kind == TK_SEMICOLON)
        token_next();
    return true;
}

decl_t* parse_decl()
{
	if (tokc.kind != TK_ID) {
        return nullptr;
    }
    decl_t *decl = new decl_t;
    if (!do_parse_decl(decl)) {
        delete decl;
        return nullptr;
    }
    return decl;
}
}

field_t* parse_field()
{
	if (tokc.kind != TK_ID) {
        return nullptr;
    }
    field_t *field = new field_t;
    if (!do_parse_decl(field)) {
        delete field;
        return nullptr;
    }
    if (tokc.kind == TK_SEMICOLON)
        token_next();
    return field;
}

static std::map<std::string, class_def_t*> g_classes;

static void complete_parent_class(class_def_t* stmt)
{
    g_classes[stmt->name] = stmt;
    int my_count = sv_count(stmt->fields);
    if (stmt->parent == nullptr) {
        stmt->field_count = sv_count(stmt->fields);
        for (int i = 0; i < my_count; ++i) {
            stmt->fields[i]->index = i;
        }
        return;
    }
    auto parent = stmt->parent;
    stmt->field_count = parent->field_count + my_count;
    for (int i = 0; i < my_count; ++i) {
        stmt->fields[i]->index = parent->field_count + i;
    }
}

class_def_t* parse_class()
{
    if (tokc.kind != TK_CLASS) {
        return nullptr;
    }
    auto* stmt = new class_def_t;
    stmt->token = tokc;
    stmt->loc = tokc.loc;
    int line = tokc.loc.line;

    // class
    NEXT_ON_KIND(stmt, TK_CLASS);
    // id: class name
    EXPECT_TOKEN(stmt, TK_ID);
    stmt->name = _strdup(ref_to_string(&tokc.ss));
    token_next();
    if (tokc.kind == TK_LPAREN) {
        token_next();
        if (tokc.kind == TK_ID) {
            auto iter = g_classes.find(ref_to_string(&tokc.ss));
            if (iter == g_classes.end()) {
                mark_line();
                do_error(&stmt->token.loc, "can't find the parent class: %s", stmt->parent->name);
                delete stmt;
                return nullptr;
            }
            stmt->parent = iter->second;
            token_next();
        }
        NEXT_ON_KIND(stmt, TK_RPAREN);
    }
    // {
    EXPECT_TOKEN(stmt, TK_LBRACE);
    token_next();
    // statements
    int field_index = 0;
    while (tokc.kind != TK_RBRACE && tokc.kind != TK_NONE) {
        while (tokc.kind == TK_SEMICOLON)
            token_next();
        if (tokc.kind == TK_DEF) {
            auto method = parse_method();
            if (method == nullptr) {
                delete stmt;
                return nullptr;
            }
            stmt->push(method);
        } else if (tokc.kind == TK_ID){
            auto field = parse_field();
            if (field == nullptr) {
                delete stmt;
                return nullptr;
            }
            field->index = field_index++;
            stmt->push(field);
        } else {
            mark_line();
            do_error(&tokc.loc, "expect a method/field/'}' here");
            delete stmt;
            return nullptr;
        }
    }
    // }
    EXPECT_TOKEN(stmt, TK_RBRACE);
    token_next();
    if (tokc.kind == TK_SEMICOLON)
        token_next();

    stmt->to_constructor();
    complete_parent_class(stmt);
    return stmt;
}

ast_node_t* parse_statement()
{
    ast_node_t* node = nullptr;

    // while (tokc.kind == TK_SEMICOLON)
    //     token_next();

    if (token_error_count() > 0) {
        return nullptr;
    }
    auto tok = tokc;
    switch (tokc.kind) {
    case TK_IF:     node = parse_if_stmt();     break;
    case TK_WHILE:  node = parse_while_stmt();  break;
    case TK_FOR:    node = parse_for_stmt();    break;
    case TK_LBRACE: node = parse_block_stmt();  break;
    case TK_IMPORT: node = parse_import_stmt(); break;
    case TK_DEF:    node = parse_method();    break;
    case TK_RETURN: node = parse_return_stmt(); break;
    case TK_GLOBAL: node = parse_global_stmt(); break;
    case TK_BREAK:  node = parse_break_stmt();  break;
    case TK_CLASS:  node = parse_class();  break;
    case TK_NONE:   node = nullptr;             break;
    case TK_SEMICOLON: node = nullptr; token_next(); break;
    default:        node = parse_expr_stmt();   break;
    }
    if (node) {
        node->loc = tok.loc;
        dbg(fprint_location(stdout, &node->loc));
        dbg(fprintf(stdout, "statement: \n"));
        dbg(print_statement(node));
        dbg(print_statement_s(node));
    }
    return node;
}

void print_statement(ast_node_t* node)
{
    do_print_statement(node, 0, false);
}

static void print_tab(int tab)
{
    for (int i = 0; i < tab; ++i)
    {
        printf("    ");
    }
}

static int get_tab(int tab, ast_node_t* node)
{
    return node->kind == NK_CompoundStatement ? tab : tab + 1;
}

static void print_expression_sel(expr_t* expr, bool s)
{
    if (s)
        print_expression_s(expr);
    else
        print_expression(expr);
}

static bool is_str(expr_t* expr)
{
    return expr->kind == NK_Expression && expr->ty == &g_string_type;
}

void do_print_statement(ast_node_t* node, int tab, bool s)
{
    switch (node->kind) {
    case NK_IfStatement: {
        if_stmt_t* stmt = (if_stmt_t*)node;
        print_tab(tab);
        printf("if (");
        print_expression_sel(stmt->cond, s);
        printf(")\n");
        int tab2 = get_tab(tab, stmt->then_stmt);
        do_print_statement(stmt->then_stmt, tab2, s);
        if (stmt->else_stmt) {
            print_tab(tab);
            printf("else\n");
            tab2 = get_tab(tab, stmt->else_stmt);
            do_print_statement(stmt->else_stmt, tab2, s);
        }
        break;
    } case NK_WhileStatement: {
        while_stmt_t* stmt = (while_stmt_t*)node;
        printf("while (");
        print_expression_sel(stmt->cond, s);
        printf(")\n");
        do_print_statement(stmt->stmt, get_tab(tab, stmt->stmt), s);
        break;
    } case NK_ForStatement: {
        for_stmt_t* stmt = (for_stmt_t*)node;
        if (stmt->init == nullptr && stmt->cond == nullptr && stmt->post == nullptr) {
            printf("for (;;)\n");
        } else {
            printf("for (\n");
            if (stmt->init) {
                do_print_statement(stmt->init, tab + 1, true);
            }
            if (stmt->cond) {
                do_print_statement(stmt->cond, tab + 1, true);
            }
            if (stmt->post) {
                do_print_statement(stmt->post, tab + 1, true);
            }
            printf(")\n");
        }
        do_print_statement(stmt->stmt, get_tab(tab, stmt->stmt), s);
        break;
    } case NK_CompoundStatement: {
        block_stmt_t* stmt = (block_stmt_t*)node;
        print_tab(tab);
        printf("{\n");
        for (int i = 0; i < sv_count(stmt->statements); ++i) {
            do_print_node(stmt->statements[i], tab + 1, s);
        }
        print_tab(tab);
        printf("}\n");
        break;
    } case NK_ExpressionStatement: {
        auto stmt = (expr_stmt_t*)node;
        print_tab(tab);
        print_expression_sel(stmt->expr, s);
        printf("\n");
        break;
    } case NK_Import: {
        auto stmt = (import_stmt_t*)node;
        print_tab(tab);
        printf("import '%s'\n", stmt->token.val.s);
        break;
    } case NK_METHOD: {
        auto stmt = (method_t*)node;
        print_tab(tab);
        printf("def %s(", stmt->name);
        for (int i = 0; i < sv_count(stmt->params); ++i) {
            if (i != 0)
                printf(", ");
            do_print_statement(stmt->params[i], 0, true);
        }
        printf(") {\n");
        for (int i = 0; i < sv_count(stmt->statements); ++i) {
            if (stmt->statements[i]->kind == NK_CompoundStatement)
                do_print_statement(stmt->statements[i], tab, true);
            else
                do_print_statement(stmt->statements[i], tab + 1, true);
        }
        print_tab(tab);
        printf("}\n");
        break;
    } case NK_ReturnStatement: {
        auto stmt = (return_stmt_t*)node;
        print_tab(tab);
        printf("return ");
        print_expression_sel(stmt->expr, s);
        printf("\n");
        break;
    } case NK_Global: {
        auto stmt = (global_stmt_t*)node;
        print_tab(tab);
        printf("global ");
        print_expression_sel(stmt->expr, s);
        printf("\n");
        break;
    } case NK_BreakStatement: {
        print_tab(tab);
        printf("break\n");
        break;
    } case NK_Class: {
        print_tab(tab);
        auto stmt = (class_def_t*)node;
        printf("class %s\n", stmt->name);
        printf("{\n");
        for (int i = 0; i < sv_count(stmt->fields); ++i) {
            do_print_statement(stmt->fields[i], tab + 1, true);
            printf("\n");
        }
        for (int i = 0; i < sv_count(stmt->methods); ++i) {
            do_print_statement(stmt->methods[i], tab + 1, true);
        }
        printf("}\n");
        break;
    } case NK_Declaration: {
        print_tab(tab);
        auto stmt = (decl_t*)node;
        print_expression_sel(stmt->decl, s);
        if (stmt->init) {
            printf(" = ");
            print_expression_sel(stmt->init, s);
        }
        break;
    } default: {
        assert(0 && "unexpected statement type");
    }
    }
}

void print_statement_ss(ast_node_t* node, int tab)
{
    switch (node->kind) {
    case NK_IfStatement: {
        if_stmt_t* stmt = (if_stmt_t*)node;
        print_tab(tab);
        printf("if (");
        print_expression_sel(stmt->cond, true);
        printf(")\n");
        break;
    } case NK_WhileStatement: {
        while_stmt_t* stmt = (while_stmt_t*)node;
        print_tab(tab);
        printf("while (");
        print_expression_sel(stmt->cond, true);
        printf(")\n");
        break;
    } case NK_ForStatement: {
        for_stmt_t* stmt = (for_stmt_t*)node;
        if (stmt->init == nullptr && stmt->cond == nullptr && stmt->post == nullptr) {
            print_tab(tab);
            printf("for (;;)\n");
        } else {
            print_tab(tab);
            printf("for (\n");
            if (stmt->init) {
                print_statement_ss(stmt->init, tab + 1);
            }
            if (stmt->cond) {
                print_statement_ss(stmt->cond, tab + 1);
            }
            if (stmt->post) {
                print_statement_ss(stmt->post, tab + 1);
            }
            printf(")\n");
        }
        break;
    } case NK_CompoundStatement: {
        block_stmt_t* stmt = (block_stmt_t*)node;
        print_tab(tab);
        printf("{\n");
        for (int i = 0; i < 1 && i < sv_count(stmt->statements); ++i) {
            print_statement_ss(stmt->statements[i], tab + 1);
        }
        print_tab(tab);
        printf("}\n");
        break;
    } case NK_ExpressionStatement: {
        auto stmt = (expr_stmt_t*)node;
        print_tab(tab);
        print_expression_sel(stmt->expr, true);
        printf("\n");
        break;
    } case NK_Import: {
        auto stmt = (import_stmt_t*)node;
        print_tab(tab);
        printf("import '%s'\n", stmt->token.val.s);
        break;
    } case NK_METHOD: {
        auto stmt = (method_t*)node;
        print_tab(tab);
        printf("def %s(", stmt->name);
        for (int i = 0; i < sv_count(stmt->params); ++i) {
            if (i != 0)
                printf(", ");
            print_statement_ss(stmt->params[i], 0);
        }
        printf("\n");
        break;
    } case NK_ReturnStatement: {
        auto stmt = (return_stmt_t*)node;
        print_tab(tab);
        printf("return ");
        print_expression_sel(stmt->expr, true);
        printf("\n");
        break;
    } case NK_Global: {
        auto stmt = (global_stmt_t*)node;
        print_tab(tab);
        printf("global ");
        print_expression_sel(stmt->expr, true);
        printf("\n");
        break;
    } case NK_BreakStatement: {
        print_tab(tab);
        printf("break\n");
        break;
    } case NK_Class: {
        print_tab(tab);
        auto stmt = (class_def_t*)node;
        printf("class %s\n", stmt->name);
        break;
    } case NK_Declaration: {
        print_tab(tab);
        auto stmt = (decl_t*)node;
        print_expression_sel(stmt->decl, true);
        if (stmt->init) {
            printf(" = ");
            print_expression_sel(stmt->init, true);
        }
        printf("\n");
        break;
    } default: {
        assert(0 && "unexpected statement type");
    }
    }
}

void print_statement_s(ast_node_t* node)
{
    do_print_statement(node, 0, true);
}

void print_expr_stmt_s(ast_node_t* node)
{
    expr_stmt_t* stmt = (expr_stmt_t*)node;
    if (stmt->kind == NK_ExpressionStatement)
        print_expression_s(stmt->expr);
}

import_stmt_t::~import_stmt_t() {
    delete ast;
    ast = nullptr;
}
